package com.itcast.stargraph.job;

import com.itcast.stargraph.api.ComfyuiApi;
import com.itcast.stargraph.pojo.ComfyuiTask;
import com.itcast.stargraph.service.impl.TextToImageServiceImpl;
import lombok.extern.slf4j.Slf4j;
import org.redisson.api.RLock;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;
import retrofit2.Call;
import retrofit2.Response;

import java.util.HashMap;

/**
 * 运行comfyui工作流的定时任务
 *
 * @Author katana
 * @Date 2025/6/12 9:26
 */
@Component
@Slf4j
public class RunWorkFlowJob {
    
    @Autowired
    private TextToImageServiceImpl textToImageService;
    
    @Autowired
    private ComfyuiApi comfyuiApi;
    
    @Autowired
    private RedissonClient redissonClient;
    
    
    @Scheduled(cron = "*/5 * * * * ?")
    public void task() {
        RLock lock = redissonClient.getLock("run:work:flow:task");
        boolean result = lock.tryLock();
        if (!result) {
            log.error("没有获取到锁,不执行");
            return;
        }
        
        try {
            //1 获取队列中优先级最高的任务
            ComfyuiTask task = textToImageService.popQueueTask();
            if (task == null) {
                return;
            }
            
            //2 添加任务到comfyui工作流中
            Call<HashMap> call = comfyuiApi.addQueueTask(task.getComfyuiRequestDto());
            Response<HashMap> response = call.execute();
            //3  处理结果
            if (!response.isSuccessful()) {
                log.error("添加任务到comfyui工作流中失败,失败信息{}", response.errorBody().string());
            }
            HashMap body = response.body();
            String promptId = String.valueOf(body.get("prompt_id"));
            task.setPromptId(promptId);
            log.info("[ComfyUI工作流任务定时执行]执行成功，任务数据：{}", task);
            
            //缓存任务ID和客户端ID的数据
            textToImageService.addStartedTask(promptId,task);
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            lock.unlock();
        }
    }
}
